"""
keras 模型相关：fit & predict
"""
from utils.plot_graph import plot_model_graph


class ModelFitPredict(object):
    """
    keras model fit & predict
    """
    def __init__(self, model, train_x, train_y, valid_x, valid_y,
                 test_x, test_y, epochs, model_name, pic_path, batch_size=72,
                 workers=8, use_multiprocessing=True):
        self.model = model
        self.train_x = train_x
        self.train_y = train_y
        self.valid_x = valid_x
        self.valid_y = valid_y
        self.test_x = test_x
        self.test_y = test_y
        self.epochs = epochs
        self.model_name = model_name
        self.pic_path = pic_path
        self.batch_size = batch_size
        self.workers = workers
        self.use_multiprocessing = use_multiprocessing

    def train_model(self):
        """

        :return:
        """
        history = self.model.fit(self.train_x, self.train_y,
                                 epochs=self.epochs,
                                 batch_size=self.batch_size,
                                 validation_data=(self.valid_x, self.valid_y),
                                 verbose=1,
                                 shuffle=False,
                                 workers=self.workers,
                                 use_multiprocessing=self.use_multiprocessing)

        loss_metrics = self.model.evaluate(self.valid_x, self.valid_y,
                                           verbose=0,
                                           batch_size=self.batch_size)

        if isinstance(loss_metrics, list):
            print('loss: {}, metrics: {}'.format(loss_metrics[0], loss_metrics[1:]))
        else:
            print('loss: {}'.format(loss_metrics))

        # plot
        plot_model_graph(model_fit=history,
                         model_name=self.model_name,
                         pic_path=self.pic_path)

    def predict(self, func=None):
        """

        :param func: 准确率的计算公式
        :return:
        """
        self.train_model()

        pred_valid = self.model.predict(self.valid_x)
        pred_test = self.model.predict(self.test_x)

        if func:
            acc = func()
            acc.update_state(y_pred=pred_valid, y_true=self.valid_y)
            valid_scores = acc.result()
            acc.reset_states()
            acc.update_state(y_pred=pred_test, y_true=self.test_y)
            test_scores = acc.result()

            print(
                '准确率得分： 验证集: {}, 测试集: {}'.format(valid_scores,
                                                 test_scores))
        return pred_valid, pred_test
